package com.newbee.config;

import com.newbee.builder.TraceNodeBuilder;
import com.newbee.model.TraceContext;
import com.newbee.model.TraceNode;
import com.newbee.model.TraceNodeStack;
import com.newbee.utils.TraceContextHolder;
import com.newbee.utils.TraceNodeStackHolder;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.scheduling.annotation.EnableAsync;
import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;

import java.util.concurrent.ThreadPoolExecutor;

/**
 * 链路追踪专用线程池：自动复制主线程的TraceContext（含traceId）
 */
@Configuration
@EnableAsync
public class TraceAsyncPoolConfig {

    @Autowired
    private TraceNodeBuilder traceNodeBuilder;

    /**
     * 核心线程池配置：
     * 1. 自定义TaskDecorator，复制主线程TraceContext到异步线程
     * 2. 合理设置线程池参数，避免资源耗尽或线程闲置
     */
    @Bean(name = "traceAsyncExecutor")
    public ThreadPoolTaskExecutor traceAsyncExecutor() {
        ThreadPoolTaskExecutor executor = new ThreadPoolTaskExecutor();

        // 1. 基础线程池参数
        // 核心线程数：默认活跃的线程数
        executor.setCorePoolSize(5);
        // 最大线程数：线程池能扩容到的最大线程数
        executor.setMaxPoolSize(10);
        // 任务队列容量：核心线程满时，暂存任务的队列大小
        executor.setQueueCapacity(20);
        // 空闲线程存活时间：非核心线程空闲60秒后销毁
        executor.setKeepAliveSeconds(60);
        // 线程名称前缀：便于日志排查
        executor.setThreadNamePrefix("TraceAsync-");

        // 2. 核心：TaskDecorator实现上下文复制
        executor.setTaskDecorator(runnable -> {
            // ① 先获取主线程的TraceContext（此时处于主线程中，可直接获取）
            TraceContext mainTraceContext = TraceContextHolder.getOrInit(true);
            TraceNode asyncNode = traceNodeBuilder.buildAsyncNode(mainTraceContext);
            // 将异步节点添加到主线程的链路中
            TraceNodeStack stack = TraceNodeStackHolder.getOrInit(null);
            stack.relatedAsyncChildNodes(asyncNode);

            // ② 包装异步任务：执行前复制上下文，执行后清空（避免线程复用污染）
            return () -> {
                // 这一步是在「异步线程」中执行的
                try {
                    // 异步任务执行前：将主线程的TraceContext复制到当前异步线程
                    TraceContext asyncTraceContext = new TraceContext(mainTraceContext.getTraceId(), asyncNode.getNodeId());

                    // 将复制后的上下文存入异步线程的ThreadLocal
                    TraceContextHolder.getTraceContext().set(asyncTraceContext);

                    // 执行真正的异步任务
                    runnable.run();

                } finally {
                    // 异步任务执行后：强制清空当前线程的TraceContext，避免线程池复用导致污染
                    TraceContextHolder.clear();
                }
            };
        });

        // 3. 任务拒绝策略：队列满+最大线程数满时，如何处理新任务（此处选择丢弃并抛异常，可按需调整）
        executor.setRejectedExecutionHandler(new ThreadPoolExecutor.AbortPolicy());

        // 4. 初始化线程池
        executor.initialize();

        return executor;
    }
}